import hqyj_mqtt
import queue
import base64
import numpy as np
import cv2
import matplotlib.pyplot as plt

def b64_to_np(image):
    # 传进来的image是一个字典，其值是一个base64格式的图像数据，需要对其进行转化为numpy数组。

    # base64.b64decode()函数将括号内的值转化为字节型（bytes）数据。
    img_data = base64.b64decode(image['image'])

    # 将字节数据转化为numpy数组
    img_np = np.frombuffer(img_data, np.uint8)

    # 使用OpenCV读取该数组
    img = cv2.imdecode(img_np, cv2.IMREAD_COLOR)

    return img

def perspective_transform(image):
    image_size = (image.shape[1], image.shape[0])
    src = np.float32(
        [[52, image_size[1]],
        [430, image_size[1]],
        [282, 105],
        [214, 105]]
    )

    dst = np.float32(
        [[image_size[0] / 4, image_size[1]],
        [image_size[0] * 3 / 4, image_size[1]],
        [image_size[0] * 3 / 4, 0],
        [image_size[0] / 4, 0]]
    )

    # 获取透视变换矩阵
    M = cv2.getPerspectiveTransform(src, dst)

    # 获取逆透视变换矩阵
    minv = cv2.getPerspectiveTransform(dst, src)

    # 调用函数进行透视变换
    image_warp = cv2.warpPerspective(image, M, image_size, flags=cv2.INTER_LINEAR)

    return image_warp,minv

def dilate_erode(image, kernel_size):
    # 对图像进行膨胀后腐蚀，iteration表示操作的次数，kernel为卷积核。
    kernel = np.ones((kernel_size, kernel_size), dtype=np.uint8)
    image_dilate = cv2.dilate(image, kernel,iterations=1)
    result_img = cv2.erode(image_dilate, kernel, iterations=1)
    return result_img


# 梯度提取车道线
def extract_line_gradient(image_warp):
    # 首先对传进来的鸟瞰图进行一次滤波
    image_Gaussian = cv2.GaussianBlur(image_warp, (5, 5), sigmaX=1)

    # 进行灰度化
    image_gray = cv2.cvtColor(image_Gaussian, cv2.COLOR_BGR2GRAY)

    # 使用sobel算子进行梯度的计算
    # 第二个参数：表示图返回图像的深度，-1表示与原图像一致
    # dx等于1表示横向计算梯度，dy等于1表示纵向计算梯度。等于0表示不在该方向计算。
    res = cv2.Sobel(image_gray, -1, dx=1, dy=0)

    # 做一个二值化
    ret, image_binary = cv2.threshold(res, 127, 255, cv2.THRESH_BINARY)

    res = dilate_erode(image_binary, 15)
    # cv2.imshow('image_binary', image_binary)
    # cv2.imshow('res', res)
    return res

# 提取白色车道线
def hlsSelect(img, thresh=(220, 255)):
    # 判断并保存白色车道线
    hls = cv2.cvtColor(img, cv2.COLOR_BGR2HLS)
    l_channel = hls[:, :, 1]
    l_channel = l_channel / np.max(l_channel) * 255
    binary_output = np.zeros_like(l_channel)
    binary_output[(l_channel > thresh[0]) & (l_channel < thresh[1])] = 1
    return binary_output

# 提取黄色车道线
def labSelect(img, thresh=(212, 220)):
    img[:, 240:, :] = (0, 0, 0)
    lab = cv2.cvtColor(img, cv2.COLOR_BGR2Lab)
    lab_b = lab[:, :, 2]
    if np.max(lab_b) > 100:
        lab_b = lab_b / np.max(lab_b) * 255
    binary_output = np.zeros_like(lab_b)
    binary_output[(lab_b > thresh[0]) & (lab_b < thresh[1])] = 1
    return binary_output


# 提取车道线
def extract_line_color(image_warp):
    # 提取白色车道线：HSL 模型，比HSV更好提取白色，只能提取白色。
    # HSL受到光照的影响比HSV弱。HSL抗干扰性比HSV强
    hls_binary = hlsSelect(image_warp)
    cv2.imshow('hls_binary', hls_binary)

    # 提取黄色车道线：lab模型。比HSV更好提取黄色，只能提取蓝、黄色
    labB_binary = labSelect(image_warp)
    cv2.imshow('labB_binary', labB_binary)

    # 将提取到的白色车道线和黄色车道线进行融合
    combined_binary = np.zeros_like(hls_binary)
    # 将在combined_binary中将hls_binary == 1或者labB_binary == 1的位置的值赋为1.
    combined_binary[(hls_binary == 1) | (labB_binary == 1)] = 1

    # 将融合后的车道线进行先膨胀后腐蚀的操作
    dilate_erode_image = dilate_erode(combined_binary, 15)
    cv2.imshow('combined_binary', dilate_erode_image)
    return dilate_erode_image

# 拟合车道线
def finding_line(dilate_erode_image):
    # 由于车是向前走的，在图中表示车是向上走的，所以通常情况下，我们更关注图像的下半部分。
    histogram = np.sum(dilate_erode_image[dilate_erode_image.shape[0] // 2:, :], axis=0)

    # 创建一个三通道图像，用来显示小窗口寻找车道线的过程。
    out_img = np.dstack((dilate_erode_image, dilate_erode_image, dilate_erode_image))

    # 获取直方图的中点位置，也就是图像的一半。
    # np.argmax()函数可以获取最大值的下标
    midpoint = histogram.shape[0] // 2

    # 得到了直方图左侧最高点的位置
    leftx_base = np.argmax(histogram[:midpoint])

    # 得到了直方图右侧最高点的位置
    rightx_base = np.argmax(histogram[midpoint:]) + midpoint

    # 获取图像中所有非零像素的x和y的位置
    # 通过numpy数组的nonzero()方法获取所有非零像素点的x和y的位置，返回的是行索引和列索引。numpy数组中的行索引对应OpenCV中的y坐标，列索引对应x坐标。
    nonzero = dilate_erode_image.nonzero()
    nonzeroy = np.array(nonzero[0])
    nonzerox = np.array(nonzero[1])

    # 定义一些小窗口的概念
    # 定义小窗口的个数
    nwindows = 10

    # 小窗口的高度
    window_height = dilate_erode_image.shape[0] // nwindows

    # 小窗口的宽度
    margin = 50

    # 小窗口内白色像素点的个数阈值
    minpix = 40

    # 初始化当前小窗口的位置，后面会持续更新。
    leftx_current = leftx_base
    rightx_current = rightx_base

    # 记录上一次的当前位置,当循环到达右半部分没有像素点的窗口的时，调到的leftx_current已经是修改过了的。
    leftx_pre = leftx_current
    rightx_pre = rightx_current

    # 创建空列表接受左侧和右侧车道线像素的索引
    left_lane_inds = []
    right_lane_inds = []

    for window in range(nwindows):
        # 计算当前窗口的上边界的y坐标
        win_y_high = dilate_erode_image.shape[0] - (window + 1) * window_height

        # 计算当前窗口的下边界的y坐标
        win_y_low = dilate_erode_image.shape[0] - window * window_height

        # 计算左边窗口的左右边界的x坐标
        win_xleft_low = leftx_current - margin
        win_xleft_high = leftx_current + margin

        # 计算右边窗口的左右边界的x坐标
        win_xright_low = rightx_current - margin
        win_xright_high = rightx_current + margin

        # 在out_img里去显示小窗口的滑动过程
        # cv2.rectangle(out_img, (win_xleft_low, win_y_high), (win_xleft_high, win_y_low), (0, 255, 0), 2)
        # cv2.rectangle(out_img, (win_xright_low, win_y_high), (win_xright_high, win_y_low), (0, 255, 0), 2)
        # cv2.imshow('out_img', out_img)

        # 找到处于窗口内的非零像素的索引
        good_left_inds = ((nonzeroy > win_y_high) & (nonzeroy < win_y_low) &
                          (nonzerox > win_xleft_low) & (nonzerox < win_xleft_high)).nonzero()[0]
        good_right_inds = ((nonzeroy > win_y_high) & (nonzeroy < win_y_low) &
                           (nonzerox > win_xright_low) & (nonzerox < win_xright_high)).nonzero()[0]

        # 将获取到的白色像素点的索引添加到列表中
        left_lane_inds.append(good_left_inds)
        right_lane_inds.append(good_right_inds)

        if len(good_left_inds) > minpix:
            leftx_current = int(np.mean(nonzerox[good_left_inds]))
        else:
            if len(good_right_inds) > minpix:
                offset = int(np.mean(nonzerox[good_right_inds])) - rightx_pre
                leftx_current = leftx_current + offset

        if len(good_right_inds) > minpix:
            rightx_current = int(np.mean(nonzerox[good_right_inds]))
        else:
            if len(good_left_inds) > minpix:
                offset = int(np.mean(nonzerox[good_left_inds])) - leftx_pre
                rightx_current = rightx_current + offset

    # 连接索引的列表，为了后续更方便的提取出这些像素点的x和y的坐标，以便进行车道的拟合。
    left_lane_inds = np.concatenate(left_lane_inds)
    right_lane_inds = np.concatenate(right_lane_inds)

    # 提取左侧和右侧车道线像素的位置
    # left_lane_inds是一个一维numpy数组，它包含了左侧车道线在滑动窗口中找到的白色像素点的x坐标的索引。
    # 通过这些索引作为所引起应用到nonzeros数组上，就可以得到相应的左侧车道线的x坐标
    # leftx包含了左侧车道线白色像素点的x的坐标
    leftx = nonzerox[left_lane_inds]
    lefty = nonzeroy[left_lane_inds]
    rightx = nonzerox[right_lane_inds]
    righty = nonzeroy[right_lane_inds]

    # 有了坐标之后，就要去对左侧和右侧车道进行多项式拟合，从而得到拟合的车道线
    # np.polyfit()，是numpy中用于进行多项式拟合的函数
    # 它接受三个参数：x，y，deg
    # x：自变量， y：因变量
    # deg：多项式次数，如果是2，则为：y = ax^2 + bx + c
    # left_fit里存放的就是a，b，c的参数
    left_fit = np.polyfit(lefty, leftx, 2)      #因为是numpy函数，所以行表示y，列表示x。
    right_fit = np.polyfit(righty, rightx, 2)

    # 使用numpy.linspace()生成一组均匀分布的数值，用于表示竖直方向上的像素坐标，方便后续车道线的绘制。
    ploty = np.linspace(0, dilate_erode_image.shape[0] - 1, dilate_erode_image.shape[0])

    # 使用多项式拟合来估计左侧和右侧车道线的x坐标
    left_fitx  =left_fit[0] * ploty ** 2 + left_fit[1] * ploty + left_fit[2]
    right_fitx = right_fit[0] * ploty ** 2 + right_fit[1] * ploty + right_fit[2]

    #计算中间车道的位置
    middle_fitx = (left_fitx + right_fitx) // 2

    # 使用不同颜色将车道线标出来
    out_img[lefty, leftx] = [255, 0, 0]
    out_img[righty, rightx] = [0, 0, 255]

    # cv2.imshow('out_img', out_img)

    return left_fitx, right_fitx, middle_fitx, ploty

def show_line(image, image_warp, dilate_erode_image, minv, left_fitx, right_fitx,  middle_fitx, ploty):
    # 目标：显示原始图像，显示透视变换后的结果，显示车道线在原图的结果,显示单纯车道线

    # 创建一个空白图像，用于在上面绘制检测到的车道线
    warp_zero = np.zeros_like(dilate_erode_image).astype("uint8")
    # 创建一个三通道的图像
    color_warp = np.dstack((warp_zero, warp_zero, warp_zero))

    # 组合车道线坐标
    # np.transpose()函数对维度进行翻转
    pts_left = np.transpose(np.vstack([left_fitx, ploty]))
    pts_right = np.transpose(np.vstack([right_fitx, ploty]))
    pts_middle = np.transpose(np.vstack([middle_fitx,ploty]))

    # 绘制车道线坐标
    cv2.polylines(color_warp, np.int32([pts_left]), isClosed=False, color=(202, 124, 0), thickness=15)
    cv2.polylines(color_warp, np.int32([pts_right]), isClosed=False, color=(202, 124, 0), thickness=15)
    cv2.polylines(color_warp, np.int32([pts_middle]), isClosed=False, color=(202, 124, 0), thickness=15)

    # 将得到的车道线的像素点根据逆透视变换映射到原始图像中
    newwarp = cv2.warpPerspective(color_warp, minv,   (image.shape[1], image.shape[0]))

    # 将逆透视变换之后的结果与原图像进行融合
    result1 = cv2.addWeighted(image, 1, newwarp, 1, 0)

    # 创建一个灰色图
    background_zero = np.zeros_like(image).astype(np.uint8) + 127

    # 创建加权融合得到
    result = cv2.addWeighted(background_zero, 1, newwarp, 1, 0)

    # cv2.imshow('result', result)
    # cv2.imshow('result1', result1)
    # cv2.imshow('image', image)
    # cv2.imshow('image_warp', image_warp)

    concatenate_image1 = np.concatenate((image, image_warp), axis=1)
    concatenate_image2 = np.concatenate((result1, result), axis=1)
    concatenate_image3 = np.concatenate((concatenate_image1, concatenate_image2), axis=0)

    cv2.imshow('concatenate_image', concatenate_image3)
    cv2.waitKey(1)

    return None






if __name__ == '__main__':

    q_mqtt_data = queue.Queue(5)


    # 1、构建MQTT客户端，连接MQTT服务器，方便和3D场景去通信，
    # 创建MQTT客户端并连接服务器。
    mqtt_client = hqyj_mqtt.MQTTClient('127.0.0.1', 21883, 'bb', 'aa', q_mqtt_data)

    while True:
        try:
            # 2、获取服务器传来的图像
            image = q_mqtt_data.get()

            # 通过关键字判断image中的数据是否为图像。
            if 'image' in image:
                # 3、将接收到的消息解析为OpenCV处理能用的格式
                image = b64_to_np(image)

                # image = cv2.imread('./1.png')

                # 用matplotlib查看图像点的坐标
                # image_plt = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                # plt.imshow(image_plt)
                # plt.show()

                # 4、对图像进行透视变换
                image_warp, minv = perspective_transform(image)
                # image_warp_copy = image_warp.copy()

                # 5、提取车道线
                # 第一种：使用梯度来提取车道线
                dilate_erode_image = extract_line_gradient(image_warp)
                # 第二种：使用颜色来提取车道线
                # extract_line_color(image_warp_copy)

                # 6、拟合车道线
                left_fitx, right_fitx, middle_fitx, ploty = finding_line(dilate_erode_image)

                # 7、将拟合后的车道线绘制到原图中。
                show_line(image, image_warp, dilate_erode_image, minv, left_fitx, right_fitx,  middle_fitx, ploty)

        except Exception as e:
            print(e)